# badminton_court_3d.py
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

class BadmintonCourt3D:
    """3D羽毛球场地绘制类 - 基于您的原始代码"""
    
    def __init__(self):
        # 标准羽毛球场地尺寸（单位：米）
        self.height = 13.4
        self.width = 6.1
        self.height_single = self.height - 0.72 * 2
        self.width_single = self.width - 0.42 * 2
        self.net_line_height = 1.98 * 2
        self.net_height = 1.55  # 网高
    
    def create_court(self, figsize=(10, 18), view_elev=20, view_azim=-45, 
                   bg_color='black', court_color='green'):
        """
        创建完整的3D羽毛球场地
        
        参数:
        figsize: 图形大小
        view_elev: 仰角
        view_azim: 方位角
        bg_color: 背景颜色
        court_color: 场地颜色
        """
        # 创建3D图形
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot(111, projection='3d')
        
        # 设置观察角度
        ax.view_init(elev=view_elev, azim=view_azim)
        
        # 绘制场地各组成部分
        self._draw_court_floor(ax, court_color)
        self._draw_doubles_lines(ax)
        self._draw_singles_lines(ax)
        self._draw_service_lines(ax)
        self._draw_net(ax)
        
        # 设置坐标轴和显示属性
        self._setup_axes(ax, bg_color)
        
        return fig, ax
    
    def _draw_court_floor(self, ax, court_color):
        """绘制场地底面"""
        x = np.array([-self.width/2-0.42, self.width/2+0.42])
        y = np.array([-self.height/2-0.72, self.height/2+0.72])
        X, Y = np.meshgrid(x, y)
        Z = np.zeros_like(X)
        ax.plot_surface(X, Y, Z, color=court_color, alpha=0.7)
    
    def _draw_doubles_lines(self, ax):
        """绘制双打场地线"""
        x = [-self.width/2, -self.width/2, self.width/2, self.width/2, -self.width/2]
        y = [-self.height/2, self.height/2, self.height/2, -self.height/2, -self.height/2]
        z = [0.01] * 5  # 稍微抬高线条，避免被底面遮挡
        ax.plot(x, y, z, 'white', linewidth=3)
    
    def _draw_singles_lines(self, ax):
        """绘制单打场地线"""
        # 左右线
        for x in [-self.width_single/2, self.width_single/2]:
            ax.plot([x, x], [-self.height/2, self.height/2], [0.01, 0.01], 'white', linewidth=3)
        
        # 上下线
        for y in [-self.height_single/2, self.height_single/2]:
            ax.plot([-self.width/2, self.width/2], [y, y], [0.01, 0.01], 'white', linewidth=3)
    
    def _draw_service_lines(self, ax):
        """绘制发球线和中线"""
        # 发球线
        for y in [-self.net_line_height/2, self.net_line_height/2]:
            ax.plot([-self.width/2, self.width/2], [y, y], [0.01, 0.01], 'white', linewidth=3)
        
        # 中线
        ax.plot([0, 0], [-self.height/2, -self.net_line_height/2], [0.01, 0.01], 'white', linewidth=3)
        ax.plot([0, 0], [self.net_line_height/2, self.height/2], [0.01, 0.01], 'white', linewidth=3)
    
    def _draw_net(self, ax):
        """绘制网"""
        # 网柱
        ax.plot([-self.width/2, -self.width/2], [0, 0], [0, self.net_height], 'black', linewidth=3)
        ax.plot([self.width/2, self.width/2], [0, 0], [0, self.net_height], 'black', linewidth=3)
        
        # 网线（用网格表示）
        x_net = np.linspace(-self.width/2, self.width/2, 50)
        z_net = np.ones_like(x_net) * self.net_height - (x_net - (-self.width/2)) * (x_net - self.width/2) * 0.0002
        ax.plot(x_net, [0]*len(x_net), z_net, 'black', linewidth=1)
        
        # 添加横向网线
        for h in np.linspace(0, self.net_height, 10):
            x_line = np.linspace(-self.width/2, self.width/2, 2)
            z_line = np.ones_like(x_line) * h - (x_line - (-self.width/2)) * (x_line - self.width/2) * 0.0002
            ax.plot(x_line, [0]*len(x_line), z_line, 'black', alpha=0.3)
    
    def _setup_axes(self, ax, bg_color):
        """设置坐标轴属性"""
        # 设置坐标轴标签和颜色
        ax.set_xlabel('Width (m)', color='white')
        ax.set_ylabel('Length (m)', color='white')
        ax.set_zlabel('Height (m)', color='white')
        
        # 设置显示范围
        ax.set_xlim(-self.width/2-1, self.width/2+1)
        ax.set_ylim(-self.height/2-1, self.height/2+1)
        ax.set_zlim(0, self.net_height*2)
        
        # 设置刻度标签颜色为白色
        ax.tick_params(axis='x', colors='white')
        ax.tick_params(axis='y', colors='white')
        ax.tick_params(axis='z', colors='white')
        
        # 设置刻度间隔
        ax.set_xticks(np.arange(-4, 5, 1))
        ax.set_yticks(np.arange(-8, 9, 2))
        ax.set_zticks(np.arange(0, 2.1, 0.5))
        
        # 设置坐标轴比例相等
        ax.set_box_aspect([1, self.height/self.width, self.net_height*2/self.width])
        
        # 设置标题
        plt.title("Badminton Court (3D)", fontsize=28, color='red', pad=20)
        
        # 显示网格
        ax.grid(True, color='gray', alpha=0.3)
        
        # 设置背景色
        ax.set_facecolor(bg_color)
        plt.gcf().patch.set_facecolor(bg_color)
    
    def add_shuttlecock_trajectory(self, ax, df, visible_col='visible', x_col='x', y_col='y', z_col='z'):
        """
        在现有场地上添加羽毛球轨迹
        
        参数:
        ax: 坐标轴对象
        df: 包含轨迹数据的DataFrame
        visible_col: 可见性列名
        x_col, y_col, z_col: 坐标列名
        """
        if visible_col in df.columns:
            df_visible = df[df[visible_col] == 1].copy()
        else:
            df_visible = df.copy()
        
        if len(df_visible) > 0:
            # 绘制轨迹点
            scatter = ax.scatter(df_visible[x_col], df_visible[y_col], df_visible[z_col], 
                               c=df_visible.index if 'frame_idx' not in df_visible.columns else df_visible['frame_idx'],
                               cmap='viridis', s=80, label='Shuttlecock', edgecolors='white')
            
            # 绘制轨迹线
            ax.plot(df_visible[x_col], df_visible[y_col], df_visible[z_col], 
                   'yellow', alpha=0.7, linewidth=2, label='Trajectory')
            
            return scatter
        return None

# 便捷函数
def create_standard_court(**kwargs):
    """创建标准羽毛球场的便捷函数"""
    court = BadmintonCourt3D()
    return court.create_court(**kwargs)

# 测试代码
if __name__ == "__main__":
    # 测试场地创建
    court = BadmintonCourt3D()
    fig, ax = court.create_court()
    plt.tight_layout()
    plt.show()